"""
简单的API接口调用性能测试

直接通过HTTP接口发送20分钟数据，测量处理时间
功能：
1. 生成20分钟的模拟数据
2. 通过API上传数据
3. 通过API获取分析结果
4. 测量总处理时间

原理：
- 使用requests库直接调用HTTP API
- 模拟真实的客户端调用场景
- 精确测量API响应时间
"""

import requests
import time
import json
import random
from datetime import datetime, timedelta
from typing import List, Dict, Any

# API配置
BASE_URL = "http://localhost:8000"
UPLOAD_API = f"{BASE_URL}/tablet/brainwave"
ANALYSIS_API = f"{BASE_URL}/tablet/realtime"

class SimpleAPITester:
    """简单的API测试器"""
    
    def __init__(self):
        """初始化测试器"""
        self.session_id = f"api_test_{int(time.time())}"
        self.room_id = "performance_test_room"
        
    def generate_mock_data_point(self, timestamp: datetime) -> Dict[str, Any]:
        """
        生成单个模拟数据点
        
        Args:
            timestamp: 时间戳
            
        Returns:
            数据点字典
        """
        # 生成512个随机EEG数据点（模拟1秒512Hz采样）
        raw_value = [random.randint(-100, 100) for _ in range(512)]
        
        # 生成频域数据（整数值）
        delta = random.randint(10, 50)
        theta = random.randint(5, 30)
        low_alpha = random.randint(5, 25)
        high_alpha = random.randint(3, 20)
        low_beta = random.randint(3, 15)
        high_beta = random.randint(2, 10)
        low_gamma = random.randint(1, 8)
        mid_gamma = random.randint(1, 5)
        
        return {
            "session_id": self.session_id,
            "room_id": self.room_id,
            "timestamp": int(timestamp.timestamp() * 1000),  # 毫秒时间戳
            "raw_value": raw_value,
            "delta": delta,
            "theta": theta,
            "low_alpha": low_alpha,
            "high_alpha": high_alpha,
            "low_beta": low_beta,
            "high_beta": high_beta,
            "low_gamma": low_gamma,
            "mid_gamma": mid_gamma,
            "attention": random.randint(20, 80),
            "meditation": random.randint(10, 70),
            "poor_signal": random.randint(0, 50),
            "blink_strength": random.randint(0, 100)
        }
    
    def generate_20min_data(self) -> List[Dict[str, Any]]:
        """
        生成20分钟的测试数据
        
        Returns:
            数据点列表
        """
        print("开始生成20分钟测试数据...")
        start_time = time.perf_counter()
        
        data_points = []
        base_timestamp = datetime.now()
        
        for i in range(1200):  # 20分钟 = 1200秒
            timestamp = base_timestamp + timedelta(seconds=i)
            data_point = self.generate_mock_data_point(timestamp)
            data_points.append(data_point)
            
            if (i + 1) % 200 == 0:
                print(f"已生成 {i + 1}/1200 个数据点...")
        
        generation_time = time.perf_counter() - start_time
        print(f"数据生成完成！用时: {generation_time:.2f}秒")
        print(f"数据大小: {len(data_points)} 个数据点")
        
        return data_points
    
    def upload_data_batch(self, data_points: List[Dict[str, Any]], batch_size: int = 50) -> float:
        """
        批量上传数据
        
        Args:
            data_points: 数据点列表
            batch_size: 批次大小
            
        Returns:
            上传总时间
        """
        print(f"\n开始上传数据，批次大小: {batch_size}")
        start_time = time.perf_counter()
        
        total_batches = (len(data_points) + batch_size - 1) // batch_size
        successful_uploads = 0
        
        for i in range(0, len(data_points), batch_size):
            batch = data_points[i:i + batch_size]
            batch_num = i // batch_size + 1
            
            try:
                # 逐个上传数据点
                for data_point in batch:
                    response = requests.post(UPLOAD_API, json=data_point, timeout=30)
                    if response.status_code == 200:
                        successful_uploads += 1
                    else:
                        print(f"上传失败: {response.status_code} - {response.text}")
                
                print(f"批次 {batch_num}/{total_batches} 完成，成功上传: {len(batch)} 个数据点")
                
            except Exception as e:
                print(f"批次 {batch_num} 上传出错: {str(e)}")
        
        upload_time = time.perf_counter() - start_time
        print(f"数据上传完成！用时: {upload_time:.2f}秒")
        print(f"成功上传: {successful_uploads}/{len(data_points)} 个数据点")
        
        return upload_time
    
    def wait_for_analysis_completion(self, max_wait_time: int = 300) -> tuple:
        """
        等待分析完成
        
        Args:
            max_wait_time: 最大等待时间（秒）
            
        Returns:
            (分析时间, 分析结果)
        """
        print(f"\n开始等待分析完成，最大等待时间: {max_wait_time}秒")
        start_time = time.perf_counter()
        
        while True:
            try:
                # 调用分析API
                params = {
                    "session_id": self.session_id,
                    "room_id": self.room_id
                }
                response = requests.get(ANALYSIS_API, params=params, timeout=10)
                
                if response.status_code == 200:
                    result = response.json()
                    status = result.get("status", "no_data")
                    
                    elapsed = time.perf_counter() - start_time
                    print(f"分析状态: {status}, 已等待: {elapsed:.1f}秒")
                    
                    if status == "completed":
                        print("分析完成！")
                        return elapsed, result
                    elif status == "error":
                        print(f"分析出错: {result.get('message', 'Unknown error')}")
                        return elapsed, result
                    elif elapsed > max_wait_time:
                        print("等待超时")
                        return elapsed, {"status": "timeout"}
                    
                    # 等待5秒后重试
                    time.sleep(5)
                else:
                    print(f"API调用失败: {response.status_code}")
                    time.sleep(5)
                    
            except Exception as e:
                print(f"查询分析状态出错: {str(e)}")
                time.sleep(5)
                
                if time.perf_counter() - start_time > max_wait_time:
                    return time.perf_counter() - start_time, {"status": "error", "message": str(e)}
    
    def run_performance_test(self):
        """运行完整的性能测试"""
        print("="*60)
        print("20分钟数据API接口性能测试")
        print("="*60)
        
        total_start_time = time.perf_counter()
        
        try:
            # 1. 生成测试数据
            data_points = self.generate_20min_data()
            
            # 2. 上传数据并测量时间
            upload_time = self.upload_data_batch(data_points)
            
            # 3. 等待分析完成并测量时间
            analysis_time, analysis_result = self.wait_for_analysis_completion()
            
            # 4. 输出性能总结
            total_time = time.perf_counter() - total_start_time
            
            print("\n" + "="*60)
            print("性能测试结果总结")
            print("="*60)
            
            print(f"数据点数量: 1200个 (20分钟)")
            print(f"数据上传时间: {upload_time:.2f}秒")
            print(f"分析等待时间: {analysis_time:.2f}秒")
            print(f"总处理时间: {total_time:.2f}秒")
            
            # 分析结果
            if analysis_result.get("status") == "completed":
                print(f"\n分析结果:")
                print(f"睡眠阶段: {analysis_result.get('sleep_stage', 'N/A')}")
                print(f"置信度: {analysis_result.get('confidence', 'N/A')}")
                print(f"数据处理数量: {analysis_result.get('data_count', 'N/A')}")
                
                # 性能评估
                if total_time < 60:
                    performance = "优秀"
                elif total_time < 180:
                    performance = "良好"
                elif total_time < 300:
                    performance = "一般"
                else:
                    performance = "需要优化"
                
                print(f"性能评级: {performance}")
            else:
                print(f"\n分析状态: {analysis_result.get('status', 'unknown')}")
                print(f"错误信息: {analysis_result.get('message', 'N/A')}")
            
            print("="*60)
            
        except Exception as e:
            print(f"测试过程中出现错误: {str(e)}")
            import traceback
            traceback.print_exc()


def main():
    """主函数"""
    # 检查服务器是否运行
    try:
        response = requests.get(f"{BASE_URL}/", timeout=5)
        print("✓ 服务器连接正常")
    except Exception as e:
        print(f"✗ 无法连接到服务器: {str(e)}")
        print("请确保服务器正在运行 (python main.py)")
        return
    
    # 运行测试
    tester = SimpleAPITester()
    tester.run_performance_test()


if __name__ == "__main__":
    main()